#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : datasets.py
@Project  : pipecoco
@Date     : 2021/8/23
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''

from mindspore import load_checkpoint, load_param_into_net
import os
from net.resnet.resnet import resnet34 as resnet
from net.resnet.resnet import ResidualBlockBase
import logging

logger = logging.getLogger("net")

def load_resnet34(model_path = None):

    network = resnet(class_num=10)
    if model_path!=None and os.path.exists(model_path):
        param_dict = load_checkpoint(model_path)
        load_param_into_net(network, param_dict)
        logger.info("{} load checkpoint from [{}].".format(network.cls_name,model_path))
    network.set_train(False)

    layer_list = []
    if network.use_se:
        for layer in ['conv1_0','bn1_0','relu','conv1_1','bn1_1','relu','conv1_2']:
            layer_list.append(getattr(network,layer))
    else:
        layer_list.append(network.conv1)
    layer_list.append(network.bn1)
    layer_list.append(network.relu)
    if network.res_base:
        layer_list.append(network.pad)
    layer_list.append(network.maxpool)
    for name in ['layer1','layer2','layer3','layer4']:
        layer = getattr(network, name)
        for block in layer:
            if isinstance(block,ResidualBlockBase):
                layer_list.extend(reconstruct_residual_base_block(block))
    layer_list.append([network.mean,(2,3)])
    layer_list.append(network.flatten)
    layer_list.append(network.end_point)
    return layer_list


'''
重构resnet的基本块
'''
def reconstruct_residual_base_block(block):

    layers_list = []
    for layer in ['conv1', 'bn1d', 'relu', 'conv2', 'bn2d']:
        layers_list.append(getattr(block, layer))
    layers_list.append('bypass')
    if block.down_sample:
        layers_list.append(block.down_sample_layer)
    else:
        layers_list.append(None)
    layers_list.append(block.relu)
    return layers_list